package net.xuele.learn.queue;

import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.atomic.AtomicStampedReference;
import java.util.function.Consumer;

/**
 * 因为kafka读取数据的时候有时候只能读取一条数据，保存这些数据在内存中的队列，
 * 多个内存的队列，一次只写一个队列，写满了就进行切换，数据写到一定的大小就执行刷盘（写ES），
 * 或者过了一段时间之后就刷盘
 * <p>
 * 这么做是为了达到批量写ES的目的
 *
 * @param <T>
 */
public class QueuePool<T> {
    /**
     * The maximum capacity, used if a higher value is implicitly specified
     * by either of the constructors with arguments.
     * MUST be a power of two <= 1<<30.
     */
    private static final int MAXIMUM_CAPACITY = 1 << 30;

    /**
     * 当前使用的queue的索引
     */
    private AtomicStampedReference<Integer> stampedReference = new AtomicStampedReference<>(0, 0);

    private SimpleCache cache = new SimpleCache(2);
    /**
     * 清除索引，最开始是0
     */
    private int cleanIndex = 0;


    private int queueSize;

    private int mask;

    /**
     * 默认实现是直接移除
     */
    private Consumer<ConcurrentLinkedQueue<T>> defaultConsumer = new Consumer<ConcurrentLinkedQueue<T>>() {
        @Override
        public void accept(ConcurrentLinkedQueue<T> t) {
            for (int i = 0; i < t.size(); i++) {
                t.remove();
            }
        }
    };

    private Consumer<ConcurrentLinkedQueue<T>> custom = null;


    private ConcurrentLinkedQueue<T>[] pool;

    public QueuePool(int poolSize, int queueSize) {
        this(poolSize, queueSize, null);
    }

    public QueuePool(int poolSize, int queueSize, Consumer<ConcurrentLinkedQueue<T>> custom) {
        this(poolSize, queueSize, 0.75F, custom);
    }


    public QueuePool(int poolSize, int queueSize, float loadFactor, Consumer<ConcurrentLinkedQueue<T>> custom) {
        poolSize = poolSize < 0 ? 2 : poolSize;
        if (queueSize > 2000) {
            throw new IllegalStateException("too large, maybe not safe for data");
        }
        if (loadFactor < 0.0F || Float.isNaN(loadFactor)) {
            throw new IllegalArgumentException("Illegal load factor: " + loadFactor);
        }
        // 将队列长度拉到最近的2的整数次幂
        this.queueSize = tableSizeFor(queueSize);
        pool = initPool(poolSize, this.queueSize, loadFactor);
        this.mask = poolSize - 1;
        this.custom = custom;
    }


    private int tableSizeFor(int cap) {
        int n = cap - 1;
        n |= n >>> 1;
        n |= n >>> 2;
        n |= n >>> 4;
        n |= n >>> 8;
        n |= n >>> 16;
        return (n < 0) ? 1 : (n >= MAXIMUM_CAPACITY) ? MAXIMUM_CAPACITY : n + 1;
    }

    private ConcurrentLinkedQueue<T>[] initPool(int poolSize, int queueSize, float loadFactor) {
        StatefulConcurrentLinkedQueue[] pool = new StatefulConcurrentLinkedQueue[poolSize];
        for (int i = 0; i < poolSize; i++) {
            pool[i] = new StatefulConcurrentLinkedQueue(queueSize, loadFactor);
        }
        return pool;
    }

    /**
     * 第一次调用此方法
     * 获取当前可写的queue，先写满其中一个（或者写到其中一个阈值），再进行切换
     * 需要考虑多线程的情况
     * <p>
     * 不同的消费者线程将同时调用此方法
     * <p>
     * 可能的情况：
     * 我们尝试take的时候，进行了切换，并且入库销毁的这个动作很快，就发生了，那么就有可两个队列同时使用了
     *
     * @return
     */
    public ConcurrentLinkedQueue<T> take() {
        // 获取当前的指针
        int pos = this.stampedReference.getReference();
        // 在数组中的索引
        int current = pos & this.mask;
        ConcurrentLinkedQueue<T> queue = pool[current];
        if (((StatefulConcurrentLinkedQueue<T>) queue).canUse()) {
            // 返回当前可用的队列，但是并不一定可用，因为其他的线程可能正在写入
            // 所以调用写入的方法时需使用不长时间阻塞的版本
            return queue;
        } else {

            // 由于当前的queue已写满了，所以这个queue中的数据就可以入库等操作了
            // 并且需要切换到下一个queue
            // 尝试去获取下一个queue，此时可能是多个线程在操作，只能允许其中的一个成功CAS
            ConcurrentLinkedQueue<T> next = switchDown();
            handleFull();
            return next;
        }
    }

    /**
     * 处理清空的逻辑
     * 异步处理
     */
    private void handleFull() {
        CompletableFuture.runAsync(() -> {
            if (custom != null) {
                custom.accept(cleanQueue());
            } else {
                defaultConsumer.accept(cleanQueue());
            }
        });
    }

    public ConcurrentLinkedQueue<T> getCurrent() {
        return pool[this.stampedReference.getReference() & this.mask];
    }

    public ConcurrentLinkedQueue<T>[] getAll() {
        return pool;
    }

    public ConcurrentLinkedQueue<T> switchDown() {
        int stamp = this.stampedReference.getStamp();
        return this.switchDown(stamp);
    }

    /**
     * 切换至下一个
     * 多个线程来获取队列，只有其中一个线程可以获取成功，
     * 失败了的，需要自旋获取或者阻塞被唤醒后再获取
     * <p>
     * 有一个问题需要解决，一批线程同时到达的时候，CAS只允许其中一个成功，
     * 但是如果有线程迟到了该怎么办呢？它会拿到最新的index去更新，刚好能够更新成功
     *
     * @return
     */
    private ConcurrentLinkedQueue<T> switchDown(int currentStamp) {
        // 当线程执行到此时，其他的线程已经完成了切换，那么就可以直接获取这个队列
        // 多个线程进来的时候，可能上一个线程已经切换完毕了，但是我拿到了最新的index，再次进行切换，就可能造成问题
        // 获取当前的指针
        Integer pos = stampedReference.getReference();
        int stamp = stampedReference.getStamp();
        // 相等才可以继续进行
        if (currentStamp < stamp) {
            // 小则已经切换成功，直接返回即可
            return pool[pos & mask];
        } else if (currentStamp > stamp) {
            throw new IllegalStateException("error stamp");
        }
        int next = pos + 1;
        // 只有一个线程可以成功
        if (stampedReference.compareAndSet(pos, next, stamp, stamp + 1)) {
            cleanIndex = pos & this.mask;
            // 还需要将之前的那个设置为不可用
            ((StatefulConcurrentLinkedQueue) pool[cleanIndex]).disable();
            // 完成切换之后设置一个状态
            cache.put(stamp, true);
            return pool[next & mask];
        } else {
            // 自旋等待完成
            for (; ; ) {
                // 判断本次是否成功
                Object o = cache.get(stamp);
                if (o != null && (Boolean) o) {
                    break;
                }
            }
            // 被唤醒，获取最新的指针，重新获取队列
            return pool[this.stampedReference.getReference() & mask];
        }
    }

    /**
     * 获得需要背清理的queue
     *
     * @return
     */
    private ConcurrentLinkedQueue<T> cleanQueue() {
        return pool[cleanIndex];
    }
}

